# 导入必要的库
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import pandas as pd

plot_real = False
plot1 = True
train = False



# 真实数据
data = {
    't1': [5.2, 0.0, 9.5, 2.3, 6.7, 3.1, 4.8, 1.0, 7.5, 0.5, 5.9, 4.2, 8.0, 2.8, 6.3,
        3.6, 4.5, 0.0, 7.2, 1.2, 5.5, 3.0, 6.9, 2.5, 6.0, 4.0, 4.7, 0.8, 7.0, 1.5],
    't2': [19.8, 23.1, 15.6, 17.4, 16.8, 18.9, 20.5, 22.3, 14.8, 24.0, 16.5, 18.2, 13.9, 19.5, 17.6,
        21.0, 18.7, 23.5, 15.3, 20.8, 17.2, 19.6, 16.4, 18.9, 17.8, 21.5, 18.0, 22.0, 14.5, 20.2],
    't3': [3.5, 6.8, 1.9, 4.2, 0.5, 5.4, 2.7, 7.2, 0.0, 3.8, 2.1, 5.6, 1.7, 4.9, 0.3,
        6.0, 2.5, 7.5, 0.0, 4.1, 2.8, 5.3, 1.6, 4.5, 0.2, 5.8, 2.3, 6.5, 0.0, 3.9],
    't4': [14.2, 10.5, 12.8, 11.5, 13.9, 15.1, 14.6, 9.7, 13.3, 16.2, 12.0, 13.5, 14.8, 11.2, 15.6,
        10.9, 13.7, 9.2, 14.0, 12.5, 13.9, 11.7, 14.5, 12.1, 15.3, 10.4, 13.2, 9.9, 13.6, 12.9]
}

df = pd.DataFrame(data)

# 统计分析
# 计算均值和方差
mean_t1 = df['t1'].mean()
var_t1 = df['t1'].var()
mean_t2 = df['t2'].mean()
var_t2 = df['t2'].var()
mean_t3 = df['t3'].mean()
var_t3 = df['t3'].var()
mean_t4 = df['t4'].mean()
var_t4 = df['t4'].var()

print('t1 均值:', mean_t1, ' 方差:', var_t1)
print('t2 均值:', mean_t2, ' 方差:', var_t2)
print('t3 均值:', mean_t3, ' 方差:', var_t3)
print('t4 均值:', mean_t4, ' 方差:', var_t4)
print("max: ", max(df['t1']), "min: ", min(df['t1']))
print("max: ", max(df['t2']), "min: ", min(df['t2']))
print("max: ", max(df['t3']), "min: ", min(df['t3']))
print("max: ", max(df['t4']), "min: ", min(df['t4']))

if plot_real:
    # 绘制直方图
    plt.hist(df['t1'], bins=10, color='orange', edgecolor='black', density=True)
    plt.title('Treadmill waiting time t1')
    plt.xlabel('Time (min)')
    plt.ylabel('frequency')
    plt.savefig('figure3_1.png')
    plt.clf()

    plt.hist(df['t2'], bins=10, color='orange', edgecolor='black', density=True)
    plt.title('Treadmill use time t2')  # 跑步机使用时间t2
    plt.xlabel('Time (min)')
    plt.ylabel('frequency')
    plt.savefig('figure3_2.png')
    plt.clf()

    plt.hist(df['t3'], bins=10, color='orange', edgecolor='black', density=True)
    plt.title('Bench press wait time t3')
    plt.xlabel('Time (min)')
    plt.ylabel('frequency')
    plt.savefig('figure3_3.png')
    plt.clf()

    plt.hist(df['t4'], bins=10, color='orange', edgecolor='black', density=True)
    plt.title('Bench press use t4')
    plt.xlabel('Time (min)')
    plt.ylabel('frequency')
    plt.savefig('figure3_4.png')
    plt.clf()



if plot1:
    # 数据仿真
    # 跑步机等待时间 t1，指数分布 λ=0.2
    t1 = np.random.exponential(scale=5, size=1000)
    t1 = t1[(t1 >= 0) & (t1 <= 9.5)]
    # 跑步机使用时间 t2，正态分布 μ=17, σ=4
    t2 = np.random.normal(loc=17, scale=4, size=1000)
    t2 = np.clip(t2, a_min=0, a_max=None)
    t2 = t2[(t2 >= 13.9) & (t2 <= 24)]

    # 卧推等待时间 t3，指数分布 λ=0.25
    t3 = np.random.exponential(scale=4, size=1000)
    t3 = t3[(t3 >= 0) & (t3 <= 7.5)]
    # 卧推使用时间 t4，正态分布 μ=13, σ=3
    t4 = np.random.normal(loc=13, scale=3, size=1000)
    t4 = np.clip(t4, a_min=0, a_max=None)
    t4 = t4[(t4 >= 9.2) & (t4 <= 16.2)]

    # 绘制直方图，生成报告中的图3.1 - 图3.4
    plt.hist(t1, bins=10, color='orange', edgecolor='black', density=True)
    plt.hist(df['t1'], bins=10, color='green', edgecolor='black', density=True, alpha=0.5)
    plt.title('Simulate treadmill waiting time t1')
    plt.xlabel('Time (min)')
    plt.ylabel('frequency')
    plt.savefig('figure4_1.png')  # 对应图3.1
    plt.clf()

    plt.hist(t2, bins=10, color='orange', edgecolor='black', density=True)
    plt.hist(df['t2'], bins=10, color='green', edgecolor='black', density=True, alpha=0.5)
    plt.title('Simulate treadmill use time t1')
    plt.xlabel('Time (min)')
    plt.ylabel('frequency')
    plt.savefig('figure4_2.png')  # 对应图3.2
    plt.clf()

    plt.hist(t3, bins=10, color='orange', edgecolor='black', density=True)
    plt.hist(df['t3'], bins=10, color='green', edgecolor='black', density=True, alpha=0.5)
    plt.title('Simulate bench press wait time t3')
    plt.xlabel('Time (min)')
    plt.ylabel('frequency')
    plt.savefig('figure4_3.png')  # 对应图3.3
    plt.clf()

    plt.hist(t4, bins=10, color='orange', edgecolor='black', density=True)
    plt.hist(df['t4'], bins=10, color='green', edgecolor='black', density=True, alpha=0.5)
    plt.title('Simulate bench press use time t3')
    plt.xlabel('Time (min)')
    plt.ylabel('frequency')
    plt.savefig('figure4_4.png')  # 对应图3.4
    plt.clf()

    # # 计算总等待时间
    # T_treadmill = t1 + t2
    # T_benchpress = t3 + t4

    # # 统计分析
    # print('t1平均值：', np.mean(t1))
    # print('跑步机总等待时间方差：', np.var(t2))
    # print('卧推总等待时间平均值：', np.mean(t3))
    # print('卧推总等待时间方差：', np.var(t4))

    print('t1 均值:', np.mean(t1), ' 方差:', np.var(t1))
    print('t2 均值:', np.mean(t2), ' 方差:', np.var(t2))
    print('t3 均值:', np.mean(t3), ' 方差:', np.var(t3))
    print('t4 均值:', np.mean(t4), ' 方差:', np.var(t4))

# BP神经网络逼近正态分布函数
x = np.linspace(0, 35, 1000)
y = (1 / (4 * np.sqrt(2 * np.pi))) * np.exp(-((x - 17) ** 2) / (2 * 4 ** 2))

# 构建神经网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.hidden = nn.Linear(1, 10)
        self.ac = nn.Tanh()
        self.output = nn.Linear(10, 1)

    def forward(self, x):
        x = self.ac(self.hidden(x))
        x = self.output(x)
        return x
    
class Net2(Net):  # 增加宽度
    def __init__(self):
        super(Net2, self).__init__()
        self.hidden = nn.Linear(1, 40)
        self.output = nn.Linear(40, 1)

    
# 定义改进后的神经网络
class Net3(nn.Module):
    def __init__(self):
        super(Net3, self).__init__()
        self.hidden1 = nn.Linear(1, 80)
        self.ac = nn.Tanh()
        self.hidden2 = nn.Linear(80,80)
        self.hidden3 = nn.Linear(80,40)
        self.output = nn.Linear(40, 1)

    def forward(self, x):
        x = self.ac(self.hidden1(x))
        x = self.ac(self.hidden2(x))
        x = self.ac(self.hidden3(x))
        x = self.output(x)
        return x
    
class Net4(nn.Module):
    def __init__(self):
        super(Net4, self).__init__()
        self.hidden1 = nn.Linear(1, 80)
        self.ac = nn.ReLU()
        self.hidden2 = nn.Linear(80,80)
        self.hidden3 = nn.Linear(80,40)
        self.output = nn.Linear(40, 1)

    def forward(self, x):
        x = self.ac(self.hidden1(x))
        x = self.ac(self.hidden2(x))
        x = self.ac(self.hidden3(x))
        x = self.output(x)
        return x

if train:
    # 数据准备
    x = np.linspace(0, 35, 1000).reshape(-1, 1)
    y = (1 / (4 * np.sqrt(2 * np.pi))) * np.exp(-((x - 17) ** 2) / (2 * 4 ** 2))

    # 数据转换为张量
    x_train = torch.from_numpy(x).float()
    y_train = torch.from_numpy(y).float()

    # 模型、损失函数和优化器
    net = Net4()
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(net.parameters(), lr=0.005)

    # 训练模型
    num_epochs = 8000
    for epoch in range(num_epochs):
        optimizer.zero_grad()
        outputs = net(x_train)
        loss = criterion(outputs, y_train)
        loss.backward()
        optimizer.step()
        if (epoch+1) % 1000 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.6f}')

    # 绘制改进后的结果
    y_pred = net(x_train).detach().numpy()
    plt.plot(x, y, label='Real function', color='orange')
    plt.plot(x, y_pred, label='neural network', color='green')  # 神经网络逼近
    plt.title('Approximation effect of BP neural network')
    plt.xlabel('time (min)')
    plt.ylabel('Probability density')
    plt.legend()
    plt.savefig('figure5_4.png')
    plt.clf()
